"""
FAQ和错误代码爬虫DAG
定时拉取阿里云数据库的产品参数，并将数据上传到dify知识库
"""
import sys
import os
from datetime import datetime
from pathlib import Path

# 将当前目录添加到Python路径
current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, current_dir)

# 🔧 重要：应用StarRocks连接补丁（必须在其他数据库相关导入之前）
try:
    from airflow_connection_patch import apply_all_patches
    patch_success = apply_all_patches()
    # print(f"🎉 StarRocks连接补丁应用结果: {patch_success}")
except Exception as e:
    print(f"❌ 导入StarRocks连接补丁失败: {str(e)}")
    # 继续执行，但记录警告
    print("⚠️ 将继续使用原始连接方式，可能会遇到超时问题")

# 导入配置的目录变量
from conn_manager import BASE_DATA_DIR, FAQ_DATA_DIR, EC_DATA_DIR

from airflow import DAG
from airflow.operators.python import PythonOperator
import logging
import traceback

logger = logging.getLogger(__name__)
logger.info(f"📂 当前DAG目录: {current_dir}")
logger.info(f"🐍 Python路径已更新，包含: {current_dir}")

# 导入自定义模块
try:
    # from faq_crawler import export_faq_by_model
    from unified_crawler import export_all_data
    from replace_all_txt import replace_characters_in_txt_files
    from upload_to_dify import upload_all_txt_files, upload_to_errorcode_dify, upload_to_faq_dify
    from file_cheker import check_file_changes
    logger.info("✅ 所有自定义模块导入成功")
except ImportError as e:
    logger.error(f"❌ 模块导入失败: {str(e)}")
    logger.error(f"🔍 详细错误信息:\n{traceback.format_exc()}")
    raise

# 配置日志
logging.basicConfig(level=logging.INFO)


def create_data_directories(**context):
    """创建数据目录结构"""
    logger.info("📁 开始创建数据目录结构...")
    try:
        # 创建基础数据目录
        base_path = Path(BASE_DATA_DIR)
        base_path.mkdir(exist_ok=True, parents=True)
        logger.info(f"✅ 基础数据目录已创建: {base_path.resolve()}")
        
        # 创建FAQ子目录
        faq_path = Path(FAQ_DATA_DIR)
        faq_path.mkdir(exist_ok=True, parents=True)
        logger.info(f"✅ FAQ数据目录已创建: {faq_path.resolve()}")
        
        # 创建错误代码子目录
        ec_path = Path(EC_DATA_DIR)
        ec_path.mkdir(exist_ok=True, parents=True)
        logger.info(f"✅ 错误代码数据目录已创建: {ec_path.resolve()}")
        
        # 验证目录是否真的存在
        for name, path in [("基础", base_path), ("FAQ", faq_path), ("错误代码", ec_path)]:
            if path.exists():
                logger.info(f"✅ {name}目录验证成功: {path.resolve()}")
            else:
                logger.error(f"❌ {name}目录创建失败: {path.resolve()}")
                raise Exception(f"{name}目录创建失败")
                
        logger.info("🎉 所有数据目录创建完成！")
        return f"数据目录结构创建成功: {BASE_DATA_DIR}"
        
    except Exception as e:
        logger.error(f"❌ 创建数据目录时出错: {str(e)}")
        raise


def run_unified_crawler(**context):
    """运行统一爬虫（同时获取FAQ、故障排除和错误代码数据）"""
    logger.info("🔧 开始运行统一数据爬虫...")
    try:
        # 添加连接预热，模拟测试脚本的成功模式
        logger.info("🔥 开始Unified连接预热...")
        from conn_manager import get_ts_ec_database_cursor
        with get_ts_ec_database_cursor() as cursor:
            cursor.execute("SELECT COUNT(*) as test_count FROM bluetti_wiki WHERE pid = '0' AND status = 1")
            test_result = cursor.fetchone()
            logger.info(f"✅ Unified连接预热成功: {test_result[0]} 条记录")
        
        # 等待一段时间，让连接完全释放
        import time
        logger.info("⏳ 等待连接释放...")
        time.sleep(3)
        
        # 执行实际的统一爬虫任务
        logger.info("🚀 开始执行实际统一爬虫任务...")
        categories = export_all_data()
        
        message = f"统一数据爬虫完成: {categories} 个类别 (包含FAQ、故障排除和错误代码数据)"
        logger.info(f"✅ {message}")
        return message
        
    except Exception as e:
        logger.error(f"❌ 统一数据爬虫运行失败: {str(e)}")
        raise

def run_faq_crawler(**context):
    """运行FAQ爬虫"""
    logger.info("🤖 开始运行FAQ爬虫...")
    try:
        # 首先进行环境检查
        logger.info("🔍 执行环境预检...")
        
        # 确保路径正确
        current_dir = os.path.dirname(os.path.abspath(__file__))
        if current_dir not in sys.path:
            sys.path.insert(0, current_dir)
        
        # 执行实际的FAQ爬虫任务
        logger.info("🚀 开始执行FAQ爬虫任务...")
        from faq_crawler import export_faq_by_model
        records, models = export_faq_by_model()
        
        message = f"FAQ爬虫完成: {records} 条记录, {models} 个型号"
        logger.info(f"✅ {message}")
        return message
        
    except Exception as e:
        logger.error(f"❌ FAQ爬虫运行失败: {str(e)}")
        logger.error(f"🔍 详细错误:\n{traceback.format_exc()}")
        raise

def run_text_cleaning(**context):
    """运行文本清理"""
    logger.info("🧹 开始运行文本清理...")
    try:
        # 直接调用函数，不传递目录参数，让其使用配置的默认目录
        processed, replacements = replace_characters_in_txt_files()
        
        message = f"文本清理完成: {processed} 个文件, {replacements} 次替换"
        logger.info(f"✅ {message}")
        return message
        
    except Exception as e:
        logger.error(f"❌ 文本清理运行失败: {str(e)}")
        raise

def run_file_checker(**context):
    """运行文件变更检测"""
    logger.info("🔍 开始运行文件变更检测...")
    try:
        # 调用文件检测函数
        check_result = check_file_changes()
        
        # 将检测结果推送到XCom，供后续任务使用
        context['task_instance'].xcom_push(key='file_check_result', value=check_result)
        
        has_changes = check_result['has_changes']
        files_count = len(check_result['files_to_upload'])
        
        message = f"文件检测完成: {'有' if has_changes else '无'}变更, {files_count} 个文件需要上传"
        logger.info(f"✅ {message}")
        return message
        
    except Exception as e:
        logger.error(f"❌ 文件检测失败: {str(e)}")
        raise

def upload_to_dify_files(**context):
    """根据文件检测结果智能上传数据到Dify知识库"""
    logger.info("☁️ 开始上传数据到Dify...")
    try:
        # 从XCom获取文件检测结果
        file_check_result = context['task_instance'].xcom_pull(key='file_check_result', task_ids='file_checker')
        
        if file_check_result and file_check_result['has_changes']:
            # 如果有文件变更，只上传变更的文件
            files_to_upload = file_check_result['files_to_upload']
            logger.info(f"🎯 检测到文件变更，将进行增量上传 {len(files_to_upload)} 个变更文件")
            
            # 使用修改后的函数进行增量上传
            try:
                # 分别调用FAQ和错误代码增量上传
                faq_results = upload_to_faq_dify(export_all=False, files_to_upload=files_to_upload)
                ec_results = upload_to_errorcode_dify(export_all=False, files_to_upload=files_to_upload)
                
                successful = faq_results['successful'] + ec_results['successful']
                total = faq_results['total'] + ec_results['total']
                message = f"Dify增量上传完成: {successful}/{total} 个变更文件上传成功"
            except Exception as e:
                logger.warning(f"⚠️ 增量上传失败: {str(e)}，回退到全量上传")
                successful, total = upload_all_txt_files()
                message = f"Dify全量上传完成: {successful}/{total} 个文件上传成功"
        else:
            # 如果没有文件变更，跳过上传
            logger.info("😌 没有检测到文件变更，跳过Dify上传")
            message = "没有文件变更，跳过上传"
            successful, total = 0, 0
        
        logger.info(f"✅ {message}")
        return message
        
    except Exception as e:
        logger.error(f"❌ Dify上传失败: {str(e)}")
        # 如果出错，尝试全量上传作为备用方案
        logger.info("🔄 尝试全量上传作为备用方案...")
        try:
            successful, total = upload_all_txt_files()
            message = f"Dify备用上传完成: {successful}/{total} 个文件上传成功"
            logger.info(f"✅ {message}")
            return message
        except Exception as backup_e:
            logger.error(f"❌ 备用上传也失败: {str(backup_e)}")
            raise

# 定义DAG
default_args = {
    'owner': 'data-team',
    'start_date': datetime(2024, 1, 1),
    'retries': 1,
}

dag = DAG(
    'unified_crawler',
    default_args=default_args,
    description='FAQ和故障排除数据爬虫工作流',
    schedule='0 2 * * 6',  # 每周六凌晨2点运行
    catchup=False,
    tags=['spider', 'knowledge_base', 'dify'],
)

# 定义任务
create_data_dir_task = PythonOperator(
    task_id='create_data_directories',
    python_callable=create_data_directories,
    dag=dag,
)

faq_crawler_task = PythonOperator(
    task_id='faq_crawler',
    python_callable=run_faq_crawler,
    dag=dag,
)

unified_crawler_task = PythonOperator(
    task_id='unified_crawler',
    python_callable=run_unified_crawler,
    dag=dag,
)

text_cleaning_task = PythonOperator(
    task_id='text_cleaning',
    python_callable=run_text_cleaning,
    dag=dag,
)

file_checker_task = PythonOperator(
    task_id='file_checker',
    python_callable=run_file_checker,
    dag=dag,
)

dify_upload_task = PythonOperator(
    task_id='dify_upload',
    python_callable=upload_to_dify_files,
    dag=dag,
)

# 定义任务依赖关系
# 基于测试脚本的成功经验，改为串行执行，并在任务间增加等待时间
# 在文本清理和Dify上传之间添加文件检测步骤
create_data_dir_task >> [unified_crawler_task, faq_crawler_task] >> text_cleaning_task >> file_checker_task >> dify_upload_task

if __name__ == "__main__":
    logger.info("🚀 DAG配置完成")
    logger.info(f"📁 使用的目录配置:")
    logger.info(f"  📂 基础数据目录: {BASE_DATA_DIR}")
    logger.info(f"  📂 FAQ数据目录: {FAQ_DATA_DIR}")
    logger.info(f"  📂 错误代码数据目录: {EC_DATA_DIR}") 
